import java.util.PriorityQueue;

public class Solution {

    class Node {
        int i, j, h;

        Node(int i, int j, int h) {
            this.i = i;
            this.j = j;
            this.h = h;
        }
    }

    public int trapRainWater(int[][] heightMap) {
        int w = heightMap[0].length;
        int h = heightMap.length;
        boolean[][] visitedMap = new boolean[h][w];
        PriorityQueue<Node> q = new PriorityQueue<>((o1, o2) -> o1.h - o2.h);
        for (int i = 0; i < h; i++) {
            q.add(new Node(i, 0, heightMap[i][0]));
            visitedMap[i][0] = true;
            q.add(new Node(i, w - 1, heightMap[i][w - 1]));
            visitedMap[i][w - 1] = true;
        }
        for (int i = 1; i < w - 1; i++) {
            q.add(new Node(0, i, heightMap[0][i]));
            visitedMap[0][i] = true;
            q.add(new Node(h - 1, i, heightMap[h - 1][i]));
            visitedMap[h - 1][i] = true;
        }

        int water = 0;
        while (!q.isEmpty()) {
            Node minNode = q.poll();

            // 上
            if (minNode.i > 0) {
                if (!visitedMap[minNode.i - 1][minNode.j]) {
                    int height = heightMap[minNode.i - 1][minNode.j];
                    if (height < minNode.h) {
                        water += (minNode.h - height);
                        height = minNode.h;
                    }
                    visitedMap[minNode.i - 1][minNode.j] = true;
                    q.add(new Node(minNode.i - 1, minNode.j, height));
                }
            }

            // 右
            if (minNode.j < w - 1) {
                if (!visitedMap[minNode.i][minNode.j + 1]) {
                    int height = heightMap[minNode.i][minNode.j + 1];
                    if (height < minNode.h) {
                        water += (minNode.h - height);
                        height = minNode.h;
                    }
                    visitedMap[minNode.i][minNode.j + 1] = true;
                    q.add(new Node(minNode.i, minNode.j + 1, height));
                }
            }

            // 下
            if (minNode.i < h - 1) {
                if (!visitedMap[minNode.i + 1][minNode.j]) {
                    int height = heightMap[minNode.i + 1][minNode.j];
                    if (height < minNode.h) {
                        water += (minNode.h - height);
                        height = minNode.h;
                    }
                    visitedMap[minNode.i + 1][minNode.j] = true;
                    q.add(new Node(minNode.i + 1, minNode.j, height));
                }
            }

            // 左
            if (minNode.j > 0) {
                if (!visitedMap[minNode.i][minNode.j - 1]) {
                    int height = heightMap[minNode.i][minNode.j - 1];
                    if (height < minNode.h) {
                        water += (minNode.h - height);
                        height = minNode.h;
                    }
                    visitedMap[minNode.i][minNode.j - 1] = true;
                    q.add(new Node(minNode.i, minNode.j - 1, height));
                }
            }
        }

        return water;
    }

    public static void main(String[] args) {
        Solution s = new Solution();
        System.out.println(s.trapRainWater(new int[][]{
                {1, 4, 3, 1, 3, 2},
                {3, 2, 1, 3, 2, 4},
                {2, 3, 3, 2, 3, 1}
        }));
    }
}